import mss
import mss.tools
from cnocr import CnOcr
from datetime import datetime
import argparse
import sys
import time

#pip install mss
def capture_screen_region(output_path="screenshot.png"):
    """
    截取桌面从(1478,0)到(1980,88)的矩形区域
    """
    try:
        with mss.mss() as sct:            # 获取所有显示器信息
            monitors = sct.monitors

            # 使用主显示器 (索引1)
            if len(monitors) < 2:
                print("未检测到主显示器")
                return None

            monitor = monitors[1]
            print(f"主显示器分辨率: {monitor['width']}x{monitor['height']}")

            # 定义截取区域坐标 (1478,0) 到 (1980,88)
            left = 1478
            top = 0
            right = 1920
            bottom = 66

            # 计算宽度和高度
            width = right - left
            height = bottom - top

            # 检查坐标是否在屏幕范围内
            if left < 0 or top < 0 or right > monitor['width'] or bottom > monitor['height']:
                print(f"错误: 截取区域超出屏幕范围")
                print(f"屏幕分辨率: {monitor['width']}x{monitor['height']}")
                print(f"截取区域: ({left},{top}) 到 ({right},{bottom})")
                return None

            if width <= 0 or height <= 0:
                print(f"错误: 截取区域宽度或高度无效")
                print(f"宽度: {width}, 高度: {height}")
                return None

            print(f"截取区域: x={left}, y={top}, width={width}, height={height}")

            # 定义截取区域
            region = {
                "top": top,
                "left": left,
                "width": width,
                "height": height
            }

            # 截取屏幕
            sct_img = sct.grab(region)

            # 保存截图
            mss.tools.to_png(sct_img.rgb, sct_img.size, output=output_path)
            print(f"截图已保存为: {output_path}")

            return output_path

    except Exception as e:
        print(f"截图失败: {str(e)}")
        return None


def ocr_recognition(image_path):
    """
    使用CnOCR对图片进行文本识别
    """
    try:
        print(datetime.now().strftime("1.%Y-%m-%d %H:%M:%S.%f")[:-3])
        print(datetime.now().strftime("2.%Y-%m-%d %H:%M:%S.%f")[:-3])

        # 初始化OCR模型
        ocr = CnOcr(det_model_name='en_PP-OCRv3_det')
        print(datetime.now().strftime("3.%Y-%m-%d %H:%M:%S.%f")[:-3])

        # 执行OCR识别
        out = ocr.ocr(image_path)
        print(datetime.now().strftime("4.%Y-%m-%d %H:%M:%S.%f")[:-3])

        print("原始识别结果:")
        print(out)

        # 筛选置信度超过45%的结果
        confidence_threshold = 0.45
        high_confidence_results = []

        for item in out:
            text = item['text']
            confidence = item['score']
            if confidence > confidence_threshold:
                high_confidence_results.append(text)

        print(f"\n筛选后（置信度>45%）的结果:")
        print(high_confidence_results)

        # 保存到字符数组
        result_array = high_confidence_results
        print(f"\n保存到字符数组的结果:")
        print(result_array)

        # 合并文本
        if result_array:
            combined_text = ' '.join(result_array)
            print(f"\n合并后的文本: {combined_text}")
        else:
            print("\n没有找到置信度超过45%的识别结果")

        for str1 in result_array:
            print('劈开的内容:' + str1)
        print(datetime.now().strftime("5.%Y-%m-%d %H:%M:%S.%f")[:-3])

        return result_array

    except Exception as e:
        print(f"OCR识别失败: {str(e)}")
        return []


def execute_ocr():
    time.sleep(10)
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='屏幕特定区域OCR识别工具')
    parser.add_argument('--output', '-o', default='screenshot.png', help='截图保存路径')
    args = parser.parse_args()

    try:
        # 1. 截取屏幕特定区域
        print("=== 开始屏幕截取 ===")
        image_path = capture_screen_region(args.output)

        if not image_path:
            print("屏幕截取失败，程序退出")
            sys.exit(1)

        # 2. 执行OCR识别
        print("\n=== 开始OCR识别 ===")
        results = ocr_recognition(image_path)

        # 3. 输出最终结果
        print("\n=== 最终识别结果 ===")
        for i, text in enumerate(results, 1):
            print(f"{i}. {text}")

    except KeyboardInterrupt:
        print("\n程序被用户中断")
    except Exception as e:
        print(f"程序执行出错: {str(e)}")


if __name__ == "__main__":
    execute_ocr()
